package com.example.tree.heap;

import java.util.*;

/**
 * 给定两个以升序排列的整形数组 nums1 和 nums2, 以及一个整数 k。
 * 定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2。
 * 找到和最小的 k 对数字 (u1,v1), (u2,v2) ... (uk,vk)。
 * <p>
 * 示例 1:
 * 输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
 * 输出: [1,2],[1,4],[1,6]
 * 解释: 返回序列中的前 3 对数：
 * [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
 * <p>
 * 示例 2:
 * <p>
 * 输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
 * 输出: [1,1],[1,1]
 * 解释: 返回序列中的前 2 对数：
 * [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
 * 示例 3:
 * <p>
 * 输入: nums1 = [1,2], nums2 = [3], k = 3
 * 输出: [1,3],[2,3]
 * 解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]
 */
public class Leetcode373_KSmallestPairs {
    public static void main(String[] args) {
//        int[] nums1 = {1,7,11};
//        int[] nums2 = {2,4,6};

        int[] nums1 = {1, 1, 2};
        int[] nums2 = {1, 2, 3};

//        int[] nums1 = {1,2};
//        int[] nums2 = {3};
        List<List<Integer>> res = kSmallestPairs3(nums1, nums2, 2);

        for (List<Integer> re : res) {
            System.out.println(re.get(0) + "," + re.get(1));
        }

    }

    /**
     * 解法三:最小堆(最优解)
     *
     * 1.初始化大小为k的最小堆，将含nums1第一个元素的前k个组合加入最小堆
     * 2.重复k次
     *   取出堆顶数组，加入结果队列
     *   如果还存在堆顶数组第一个元素的组合，则将下一个组合加入最小堆
     *
     * @param nums1
     * @param nums2
     * @param k
     * @return
     */

    public static List<List<Integer>> kSmallestPairs3(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new LinkedList<>();
        if (nums1.length == 0 || nums2.length == 0)
            return result;// 两个数组有一个为空，返回空

        Queue<int[]> minHeap = new PriorityQueue<>(k,
                (o1, o2) -> (nums1[o1[0]] + nums2[o1[1]]) - (nums1[o2[0]] + nums2[o2[1]]));
        // 维护一个大小为k的小顶堆，将某个升序序列加入
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            minHeap.add(new int[]{i, 0}); // 注意加入的是坐标，这样有利于得到堆顶元素的下一个元素
        }

        while (k > 0 && !minHeap.isEmpty()) {
            int[] pair = minHeap.poll(); // 弹出堆顶元素
            List<Integer> item = new ArrayList<>();
            item.add(nums1[pair[0]]);
            item.add(nums2[pair[1]]);
            result.add(item);
            if (pair[1] < nums2.length - 1) {
                minHeap.offer(new int[]{pair[0], pair[1] + 1});// 加入堆顶元素的下一个元素
            }
            k--;
        }
        return result;
    }

    /**
     * 解法二:最大堆
     * 由于两个数组都是从小到大排序的，那么和最小的 k 对数字必定在两个数组的前k个元素组成的数字对中
     * 先将两个数组前k个元素两两组合的所有可能全部列出，然后构建一个容量为k的最大堆
     *
     * 当最大堆未满或者当前数组和比堆顶的数组和小，才将当前的数组对加入最大堆
     *   如果最大堆容量已经为K了，当遇到比堆顶小的元素要先将堆顶出堆，然后再将该元素入堆
     *
     * @param nums1
     * @param nums2
     * @param k
     * @return
     */
    public static List<List<Integer>> kSmallestPairs2(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new ArrayList<>();
        if (nums1.length == 0 || nums2.length == 0)
            return result;

        Queue<Pair> maxHeap = new PriorityQueue<>(k,
                (o1, o2) -> o2.num2 + o2.num1 - (o1.num2 + o1.num1));

        for(int i = 0; i < Math.min(k, nums1.length); i++) {
            for (int j = 0; j < Math.min(k, nums2.length); j++) {
                Pair curPair = new Pair(nums1[i], nums2[j]);
                // 如果最大堆未满或者当前数组和比堆顶的数组和小，才将当前的数组对加入最大堆
                if (maxHeap.size() != k || maxHeap.peek().compareTo(curPair) >= 0) {
                    if (maxHeap.size() == k) { // 如果堆中的元素已经有了k个,则需要先将堆顶出堆
                        maxHeap.poll();
                    }
                    maxHeap.offer(curPair);
                }
            }
        }

        while (!maxHeap.isEmpty()) {
            Pair pair = maxHeap.poll();
            List<Integer> e = new ArrayList<>();
            e.add(pair.num1);
            e.add(pair.num2);
            result.add(e);
        }
        return result;
    }

    /**
     * 解法一:暴力法
     * <p>
     * 由于两个数组都是从小到大排序的，那么和最小的 k 对数字必定在两个数组的前k个元素组成的数字对中
     * 先将两个数组前k个元素两两组合的所有可能全部列出，然后对和进行排序，最后取前k个
     *
     * @param nums1
     * @param nums2
     * @param k
     * @return
     */
    public static List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new ArrayList<>();
        if (nums1.length == 0 || nums2.length == 0)
            return result;
        Pair[] pairs = new Pair[Math.min(k, nums1.length) * Math.min(k, nums2.length)];

        int index = 0;
        for (int i = 0; i < Math.min(k, nums1.length); i++) {

            for (int j = 0; j < Math.min(k, nums2.length); j++) {
                pairs[index++] = new Pair(nums1[i], nums2[j]);
            }
        }

//        Arrays.sort(pairs, new Comparator<Pair>() {
//            @Override
//            public int compare(Pair o1, Pair o2) {
//                return o1.num2 + o1.num1 - (o2.num2 + o2.num1);
//            }
//        });

        Arrays.sort(pairs);

        for (int i = 0; i < Math.min(k, pairs.length); i++) {
            List<Integer> e = new ArrayList<>();
            e.add(pairs[i].num1);
            e.add(pairs[i].num2);
            result.add(e);
        }

        return result;

    }


    private static class Pair implements Comparable<Pair>{
        int num1;
        int num2;

        public Pair(int num1, int num2) {
            this.num1 = num1;
            this.num2 = num2;
        }

        @Override
        public int compareTo(Pair other) {
            return this.num2 + this.num1 - (other.num2 + other.num1);
//            return other.num2 + other.num1 - (this.num2 + this.num1);
        }
    }
}
